Skip to content

Adding readManyByPartitionKey API#48801

Merged
FabianMeiswinkel merged 89 commits into
mainfrom
users/fabianm/readManyByPK
Apr 24, 2026
Merged

Adding readManyByPartitionKey API#48801
FabianMeiswinkel merged 89 commits into
mainfrom
users/fabianm/readManyByPK

Conversation

@FabianMeiswinkel

@FabianMeiswinkel FabianMeiswinkel commented Apr 13, 2026

Copy link
Copy Markdown
Member

Description

Adds a new readManyByPartitionKey API surface to the Java Cosmos SDK (sync + async) and wires it through the Spark connector to support PK-only reads (including partial HPK), with query-plan-based validation for custom queries.

Changes:

Added public readManyByPartitionKeys overloads in CosmosAsyncContainer / CosmosContainer and an internal AsyncDocumentClient + RxDocumentClientImpl implementation that groups PKs by physical partition and issues per-range queries.
Introduced ReadManyByPartitionKeyQueryHelper to compose PK filters into user-provided SQL and added a new config knob for per-partition batching.
Added Spark support (UDF + PK serialization/parsing helper + reader) and unit/integration tests for query composition and end-to-end behavior.

All SDK Contribution checklist:

  • The pull request does not introduce [breaking changes]
  • CHANGELOG is updated for new features, bug fixes or other significant changes.
  • I have read the contribution guidelines.

General Guidelines and Best Practices

  • Title of the pull request is clear and informative.
  • There are a small number of commits, each of which have an informative message. This means that previously merged commits do not appear in the history of the PR. For more information on cleaning up the commits in your PR, see this page.

Testing Guidelines

  • Pull request includes test coverage for the included changes.

Copilot AI review requested due to automatic review settings April 13, 2026 23:00
@FabianMeiswinkel FabianMeiswinkel marked this pull request as draft April 13, 2026 23:01

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a new readManyByPartitionKey API surface to the Java Cosmos SDK (sync + async) and wires it through the Spark connector to support PK-only reads (including partial HPK), with query-plan-based validation for custom queries.

Changes:

  • Added public readManyByPartitionKey overloads in CosmosAsyncContainer / CosmosContainer and an internal AsyncDocumentClient + RxDocumentClientImpl implementation that groups PKs by physical partition and issues per-range queries.
  • Introduced ReadManyByPartitionKeyQueryHelper to compose PK filters into user-provided SQL and added a new config knob for per-partition batching.
  • Added Spark support (UDF + PK serialization/parsing helper + reader) and unit/integration tests for query composition and end-to-end behavior.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 11 comments.

Show a summary per file
File Description
sdk/cosmos/docs/readManyByPartitionKey-design.md Design doc describing the new API, query validation, and Spark integration approach.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/query/DocumentQueryExecutionContextFactory.java Adds a helper method to fetch query plans through the gateway for validation.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/RxDocumentClientImpl.java Implements readManyByPartitionKey execution, validation, PK→range grouping, batching, and concurrency.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelper.java New helper to build SqlQuerySpec by appending PK filters and extracting table aliases.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/Configs.java Adds config/env accessors for max PKs per per-partition query batch.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/AsyncDocumentClient.java Adds internal interface method for PK-only read-many.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosContainer.java Adds sync readManyByPartitionKey overloads.
sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java Adds async readManyByPartitionKey overloads and wiring to internal client.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/implementation/ReadManyByPartitionKeyQueryHelperTest.java Unit tests for SQL generation, alias extraction, and WHERE detection.
sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ReadManyByPartitionKeyTest.java Emulator integration tests for single PK + HPK, partial HPK, projections, and query validation.
sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKeyITest.scala Spark integration test for reading by PKs and empty result behavior.
sdk/cosmos/azure-cosmos-spark_3/src/test/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelperSpec.scala Unit tests for PK string serialization/parsing helpers.
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/udf/GetCosmosPartitionKeyValue.scala Spark UDF to compute _partitionKeyIdentity values.
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala Spark partition reader that calls new SDK API and converts results to rows.
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosReadManyByPartitionKeyReader.scala Spark reader that maps input rows to PKs and streams results via the partition reader.
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosPartitionKeyHelper.scala Helper for PK serialization/parsing used by the UDF and data source.
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosItemsDataSource.scala Adds Spark entry point to read-many by partition key, including PK extraction logic.
sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/CosmosConstants.scala Adds _partitionKeyIdentity constant.
Comments suppressed due to low confidence (1)

sdk/cosmos/azure-cosmos-spark_3/src/main/scala/com/azure/cosmos/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala:1

  • The error message has mismatched parentheses/quoting (classOf<SparkRowItem])) which makes it harder to read and search for. Suggest correcting it to a clean, unambiguous string (e.g., classOf[SparkRowItem]) to improve diagnosability.
// Copyright (c) Microsoft Corporation. All rights reserved.

Comment thread sdk/cosmos/docs/readManyByPartitionKey-design.md
FabianMeiswinkel and others added 18 commits April 14, 2026 16:49
…s/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…s/spark/ItemsPartitionReaderWithReadManyByPartitionKey.scala

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…ntation/RxDocumentClientImpl.java

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
…ntation/ReadManyByPartitionKeyQueryHelper.java

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
@FabianMeiswinkel FabianMeiswinkel marked this pull request as ready for review April 16, 2026 21:57
@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

@sdkReviewAgent

@xinlian12 xinlian12 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks~

@xinlian12

Copy link
Copy Markdown
Member

Review complete (41:15)

Posted 2 inline comment(s).

Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

/azp run java - cosmos - spark

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

/azp run java - cosmos - tests

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

@sdkReviewAgent

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@xinlian12 xinlian12 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks

@xinlian12

Copy link
Copy Markdown
Member

Review complete (33:55)

Posted 2 inline comment(s).

Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

@sdkReviewAgent

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

/azp run java - cosmos - tests

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

/azp run java - cosmos - spark

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

1 similar comment
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 1 pipeline(s).

@xinlian12 xinlian12 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks

@xinlian12

Copy link
Copy Markdown
Member

Review complete (32:34)

Posted 2 inline comment(s).

Steps: ✓ context, correctness, cross-sdk, design, history, past-prs, synthesis, test-coverage

@FabianMeiswinkel

Copy link
Copy Markdown
Member Author

Live test failures independent of the changes in this PR

@FabianMeiswinkel FabianMeiswinkel merged commit e735102 into main Apr 24, 2026
108 of 110 checks passed
@FabianMeiswinkel FabianMeiswinkel deleted the users/fabianm/readManyByPK branch April 24, 2026 06:33
jeet1995 added a commit to jeet1995/azure-sdk-for-java that referenced this pull request Jun 2, 2026
…public

The upstream ReadManyByPartitionKeyQueryPlanValidationTest (added in Azure#48801)
constructs PartitionedQueryExecutionInfo directly from a test package.
Upstream/main exposed this ctor as public; restore that visibility so the
new test compiles after the merge.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
jeet1995 added a commit that referenced this pull request Jul 9, 2026
…V2 endpoint (#47759)

* Route Execute Stored Procedure requests to Thin Proxy endpoint.

* Route Execute Stored Procedure and QueryPlan requests to Thin Proxy endpoint.

* Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC + GW Mode) cases.

* Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC + GW Mode) cases.

* Ensure QueryPlan gets routed to Gateway Service Endpoint (in non-TC + GW Mode) cases.

* Obtain List<Range<String>> from List<PartitionKeyInternal>.

* Obtain List<Range<String>> from List<PartitionKeyInternal>.

* Adding query + thin-client tests.

* Fixing tests.

* Fixing tests.

* Fixing tests.

* Addressing review comments.

* Addressing review comments.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Refactor thin-client E2E tests based on operation type.

* Add SupportedQueryFeatures and QueryVersion RNTBD request headers for QueryPlan proxy routing

Add RNTBD token mappings for x-ms-cosmos-supported-query-features (0x002B)
and x-ms-cosmos-query-version (0x002C) so the thin client proxy can read
these values from the RNTBD body when processing QueryPlan requests.

Previously these headers were only set as HTTP headers by QueryPlanRetriever
and were lost when QueryPlan was routed through the proxy path, since
ThinClientStoreModel serializes requests as RNTBD (not HTTP headers).

IDs match server-side proxy definitions per ADO PR 1982503.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add change feed tests for FeedRange.forFullRange and forLogicalPartition

Add testThinClientChangeFeedFullRange covering FeedRange.forFullRange()
across multiple partition keys, and testThinClientChangeFeedPartitionKey
covering FeedRange.forLogicalPartition with exact doc count + PK validation.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add thin client E2E test matrix documentation for QueryPlan review

Documents all 59 thin client E2E tests across query (50), point operations (3),
change feed (3), and stored procedures (3) with SQL, query features covered,
and known account-side blockers.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add SupportedQueryFeatures and QueryVersion RNTBD request headers for QueryPlan proxy routing

Add RNTBD token mappings for x-ms-cosmos-supported-query-features (0x00F0)
and x-ms-cosmos-query-version (0x00F1) so the thin client proxy can read
these values from the RNTBD body when processing QueryPlan requests.

IDs are provisional (0x00F0, 0x00F1) — must be coordinated with server-side
proxy team. See ADO PR 1982503 for the proxy-side design.

Note: The design doc listed 0x002B/0x002C but those are already assigned to
PartitionKey/PartitionKeyRangeId in the Java SDK. Using 0x00F0/0x00F1 to
avoid ID collision until final server-side IDs are assigned.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix testGetCurrentDateTime flaky assertion, add AAD auth support, RNTBD instructions

- Fix testGetCurrentDateTime: assert ISO 8601 format instead of exact match
  (gateway and proxy return slightly different timestamps)
- Add DefaultAzureCredential support via COSMOS.USE_AAD_AUTH system property
  for accounts with disableLocalAuth=true
- Add RNTBD class reference as .github/instructions/rntbd.instructions.md
- Add pom.xml system properties for THINCLIENT_ENABLED, HTTP2_ENABLED, USE_AAD_AUTH
- Add beforeSuiteReuse mode for degraded accounts

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Refactor thin client query tests for reliability

- Switch baseline from Gateway V1 to Direct TCP to avoid JVM config
  interference (THINCLIENT_ENABLED/HTTP2_ENABLED affect Gateway V1)
- Assert :10250 endpoint only on Gateway V2 results (not baseline)
- Rename helpers: assertDirectAndThinClientMatch (was gateway)
- Document seedTestData schema in Javadoc
- Remove 'Expected to fail' comments (account has vector search enabled)
- Clean up class/method Javadoc

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix container leaks and Direct TCP AAD auth in tests

- Fix container leak: get container reference before createContainer() so
  finally block can always delete. Use safeDeleteContainer() helper.
- Fix Direct TCP client: apply AAD credential via applyCredential() for
  accounts with disableLocalAuth=true.

All 89 thin client tests pass (80 query + 3 change feed + 3 point ops + 3 sprocs).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Use bulkDelete in AfterClass for seeded docs cleanup

Replace sequential deleteItem loop with executeBulkOperations for
faster and more reliable cleanup of seeded test documents.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR review feedback

1. Remove rntbd.instructions.md (not part of PR)
2. Move MAPPER variable to top of PartitionKeyInternalTest
3. Remove COSMOS.REUSE_DATABASE_ID (beforeSuiteReuse, afterSuite guard, pom.xml)
4. Keep System.setProperty in ThinClientQueryE2ETest (required — extends
   TestSuiteBase not ThinClientTestBase, property must be set before client build)
5. Fix RNTBD IDs to match server-side RntbdConstants.cs (ADO PR 1982503):
   SupportedQueryFeatures=0x00FF (String), QueryVersion=0x0100 (SmallString)
   The SmallString type was the root cause of 400 BadRequest — String uses
   2-byte length prefix, SmallString uses 1-byte, causing token stream misparse.

All 89 thin client tests pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Align vector/FTS/hybrid tests with existing patterns

- Add euclidean VectorDistance variant to vector search test
- Add document ID comparison to FTS and hybrid tests (was count-only)
- Add testFullTextScoreRanking: ORDER BY RANK FullTextScore with exact
  order comparison between Direct and thin client
- All assertions now validate both result size AND content parity

90/90 thin client tests pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review agent comments + fix broken link

Code fixes:
- PartitionKeyInternalHelper: materialize fieldNames() Iterator in error msg
- QueryPlanRetriever: throw IllegalStateException when partitionKeyDef null
  in thin client mode (was silent fallthrough)
- RxDocumentClientImpl: add parentheses around compound && in boolean exprs
- ThinClientStoreModel: use != instead of !(==)
- ThinClientTestBase: remove stale 'uncomment' comment
- ThinClientQueryE2ETest: accept status 0 (transport rejection) for invalid
  query — proxy returns transport error, not HTTP 400
- THINCLIENT_TEST_MATRIX.md: remove broken link, update methodology to
  reflect Direct TCP baseline

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review agent comments (round 2)

- RntbdOperationType: add QueryPlan case to fromId()/fromType() lookups
- PartitionKeyInternalHelper: sanitize PII from error message (log node type
  instead of raw partition key values)
- THINCLIENT_TEST_MATRIX.md: remove PR-specific references and known blockers
- PartitionedQueryExecutionInfo: remove stale queryRanges from backing
  ObjectNode after EPK conversion to prevent data inconsistency
- QueryPlanRetriever: add TODO for 3 missing query features vs .NET SDK
  (ListAndSetAggregate, CountIf, HybridSearchSkipOrderByRewrite)

90/90 thin client tests pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Refactor queryRanges deserialization — detect format from response, no ObjectNode mutation

- PartitionedQueryExecutionInfo: getQueryRanges() now inspects the first
  range's 'min' field to detect format (string=EPK hex, array=PK-internal)
  and applies the correct deserialization path automatically.
- Remove ObjectNode.remove('queryRanges') — no longer mutate the backing
  JSON. The format is detected lazily on first access.
- 3-arg constructor now takes PartitionKeyDefinition (not pre-computed ranges)
  so conversion happens inside getQueryRanges(), keeping the DTO self-contained.
- QueryPlanRetriever: pass partitionKeyDefinition to constructor instead of
  pre-converting ranges externally.

90/90 thin client tests pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Use agnostic QueryRangesFormat hint instead of PartitionKeyDefinition presence

Introduce QueryRangesFormat enum (EPK_HEX_STRING, PARTITION_KEY_INTERNAL_ARRAY)
as the deserialization hint. The hint guides which path to try first, but
getQueryRanges() always validates by inspecting the actual min node type —
if the hint doesn't match the wire format, it falls through to the other path.

Naming is transport-agnostic (no Proxy/ThinClient/Gateway references).
PartitionKeyDefinition is only stored when needed for EPK conversion,
not used as a format signal.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix fetchQueryPlanForValidation caller after merge

Pass null for the DocumentCollection parameter so the public
fetchQueryPlanForValidation entry point compiles against the new
fetchQueryPlan signature introduced during the merge with upstream/main.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Make PartitionedQueryExecutionInfo(ObjectNode, RequestTimeline) ctor public

The upstream ReadManyByPartitionKeyQueryPlanValidationTest (added in #48801)
constructs PartitionedQueryExecutionInfo directly from a test package.
Upstream/main exposed this ctor as public; restore that visibility so the
new test compiles after the merge.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* [Cosmos] Default thin-client to enabled and add HTTP/2 connectivity-probe gate

Adds an EndpointOrchestrator that fans out POST /connectivity-probe to
every thin-client regional endpoint after each topology refresh. SDK
only routes data-plane traffic to thin-client (Gateway V2) when all
regional probes succeed across N consecutive refresh cycles
(configurable via COSMOS.THINCLIENT_PROBE_FAILURE_THRESHOLD, default 2);
otherwise traffic falls back to Gateway V1 at the next refresh
boundary. No mid-flight fallback.

Caveats:
- Probe wiring is skipped entirely for Direct mode and when HTTP/2 is
  not configured; controlled by RxDocumentClientImpl.useThinClient.
- QueryPlan, metadata reads, and AllVersionsAndDeletes change feed
  continue to route through Compute Gateway (Gateway V1).
- Probe failures are absorbed inside the orchestrator and the trigger
  is fire-and-forget on the GEM scheduler, so probe issues can never
  trip CosmosClient initialization or fail a topology refresh.
- EndpointOrchestrator implements Closeable and is closed from
  GlobalEndpointManager.close() so no further probes are issued after
  client shutdown.

THINCLIENT_ENABLED now defaults to true; opt out via
COSMOS.THINCLIENT_ENABLED=false or COSMOS.THINCLIENT_PROBE_ENABLED=false.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR #49437 review comments: probe single-flight, body-drain lifecycle, cancellable in-flight probes

- EndpointOrchestrator: fold body-drain into probe Mono via flatMap+then(perProbeTimeout) so a slow/hanging response body cannot leak resources outside the cycle budget (Copilot #1, deep-review #3, jeet HIGH-2 minor).

- EndpointOrchestrator: add single-flight CAS (cycleInProgress) plus monotonic cycle id; closed-check inside applyCycleResult drops late results so a post-close cycle cannot mutate health state (deep-review #1+#2, jeet HIGH-2).

- EndpointOrchestrator: re-evaluate closed/feature-flag/endpoints at subscription time via Mono.defer so GEM.close() cancellation is honored before any HTTP I/O is issued.

- GlobalEndpointManager: retain probe Disposable in AtomicReference; close() now disposes the in-flight probe subscription so probe work cannot outlive the GEM/CosmosClient (Copilot #2, deep-review #2).

- CHANGELOG: moved entry to unreleased 4.82.0-beta.1, reworded to honestly describe optimistic startup, N=2 RED-to-fallback hysteresis, and Direct-mode/metadata exclusion (Copilot #3, deep-review #4).

Tests: 45 unit tests pass (EndpointOrchestratorTests + ConfigsTests + ThinClientProbeWiringTests).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR #49437 second-batch review feedback

- LocationCache.getThinClientRegionalEndpoints now walks both read and write region endpoint maps so single-master write-region failures still flip the probe gate.
- EndpointOrchestrator.forceUnhealthy(reason) provides a non-HTTP path to flip the gate; GlobalEndpointManager calls it when topology says thin-client is eligible but no regional endpoint resolves.
- Symmetric hysteresis: new COSMOS.THINCLIENT_PROBE_RECOVERY_THRESHOLD (default 1) so operators can require N consecutive GREEN cycles before flipping back to proxy.
- Extracted RxDocumentClientImpl.useThinClientStoreModel(...) body into package-private static shouldUseThinClientStoreModel for direct unit testability; added ThinClientRoutingGateTests covering 9 routing paths.
- EndpointOrchestratorTests.stubResponse now returns Mono.empty() to avoid Unpooled.EMPTY_BUFFER refCnt underflow across multiple probe calls.
- Removed unused locals; added recoveryThresholdRequiresMultipleGreenCycles, forceUnhealthy_flipsGateToRedWithoutRunningProbe, forceUnhealthy_onClosedOrchestrator_isNoOp tests.

All 57 unit tests in the touched files pass.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Clarify CHANGELOG: probe recovery threshold is now configurable

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR #49437 review (rounds 4-8): reactor-chain probe, FQN cleanup, fix gwV2Cto and ThinClient user-agent assertions

- GlobalEndpointManager: convert thin-client probe trigger to a Mono<Void>
  chained into the topology-refresh reactor pipeline (replaces fire-and-forget
  subscribe). Removes thinClientProbeDisposable field and its close() handling
  since cancellation now propagates through the outer subscription.
- EndpointProbeClient/EndpointProbeClientTests/ThinClientProbeWiringTests:
  replace inline FQNs with imports (java.io.Closeable, java.util.List,
  java.net.ConnectException, com.azure.cosmos.implementation.http.HttpHeaders).
- ClientConfigDiagnosticsTest: compute gwV2Cto dynamically from
  Configs.isThinClientEnabled() so assertions remain valid after the default
  flip to true.
- ConfigsTests: update default-threshold assertions from 2 to 1 to match
  DEFAULT_THINCLIENT_PROBE_FAILURE_THRESHOLD=1.
- UserAgentContainerTest.UserAgentIntegration: expect '|F4' suffix because
  the ThinClient feature flag (1 << 2) is now included by default.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix CI test fallout from default ThinClient enablement

UserAgentSuffixTest.validateUserAgentSuffix and
CosmosDiagnosticsTest.generateHttp2OptedInUserAgentIfRequired:
include UserAgentFeatureFlags.ThinClient in computed |F<hex> suffix
when COSMOS.THINCLIENT_ENABLED is true (now default after Gateway V2
default enablement). Mirrors RxDocumentClientImpl.addUserAgentSuffix +
UserAgentContainer.setFeatureEnabledFlagsAsSuffix behavior.

SinglePartitionDocumentQueryTest.querySinglePartitionDocuments:
spy on both gateway-proxy and thin-proxy and assert exactly one
invocation. Previous code only spied on the proxy implied by
useThinClient() config intent, which races with the probe-healthy
gate -- routing AND's intent with isProxyProbeHealthy() so on first
cycle the request may go through gateway even when thin-client is
configured.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Disable thin-client probe in Http2PingKeepaliveTest

The test installs an iptables DROP on thin-client port 10250 to verify
that Http2PingHandler closes the broken connection after consecutive
PING ACK timeouts and the recovery request uses a new connection on
the same regional endpoint.

After default Gateway V2 enablement, the connectivity probe also fires
HTTP/2 POSTs to port 10250 on every account refresh. With iptables
dropping that port, the probe trips proxyHealthy=false, useThinClient
StoreModel() returns false, and the data plane request routes through
Gateway V1 on port 443 -- which iptables is not dropping. Result:
the PING handler never fires, the warm-up and recovery requests use
the same gateway channel, and the assertion 'recovery channel must
differ from initial' fails (both ended up as 77af2e47 on build 6419227).

Set COSMOS.THINCLIENT_PROBE_ENABLED=false in beforeClass so the probe
short-circuits to a no-op, EndpointProbeClient.proxyHealthy stays
optimistically true, and the data plane request actually flows over
port 10250 where the iptables DROP can take effect.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Disable thin-client probe by default for E2E tests in TestSuiteBase

Build 6424287 surfaced two new failure patterns:

  1. CosmosNotFoundTests.performBulkOnDeletedContainerWithGatewayV2 (45
     failures) - asserts substatus 1003 from the thin-client routing
     path, but observed 0 because the data plane was routed to Gateway
     V1 instead of the proxy.
  2. PerPartitionCircuitBreakerE2ETests.*Gateway (26 failures) -
     TestSuiteBase.assertThinClientEndpointUsed could not find any
     request whose endpoint contained ':10250/', i.e. nothing actually
     went to the thin-client proxy.

Both patterns trace to the same source: the new connectivity probe is
enabled by default, the proxy-side /connectivity-probe endpoint is not
deployed in every CI test account yet, and the default failure
threshold is 1. So after the first probe cycle the SDK marks the proxy
unhealthy and routes data plane traffic to Gateway V1, which breaks
tests that explicitly assert thin-client routing.

Disable the probe by default in TestSuiteBase's static initializer
(only when the property is not already set), so all E2E tests inherit
deterministic, configuration-driven routing. Tests that exercise the
probe itself (EndpointProbeClientTests, ThinClientProbeWiringTests)
set the property explicitly in @BeforeMethod and are not affected.

Also drop the now-redundant per-class override in Http2PingKeepaliveTest
- the base class disables it, and the test's @afterclass clear would
otherwise re-enable the probe for any subsequent E2E test sharing the
JVM.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Revert "Disable thin-client probe by default for E2E tests in TestSuiteBase"

This reverts commit ad9e3df.

* Add per-class thinclient probe disable in CosmosNotFoundTests and PerPartitionCircuitBreakerE2ETests

Companion to the prior revert. The revert undid the global TestSuiteBase probe

disable (which masked production behaviour). This commit adds the necessary

per-class disable to the two test classes whose assertions explicitly require

thinclient routing: CosmosNotFoundTests (thinclient group) and

PerPartitionCircuitBreakerE2ETests (fi-thinclient-multi-master group). Both

clear the property in their @afterclass. Http2PingKeepaliveTest already has

its own disable (restored by the revert). Production callers continue to get

the connectivity probe ON by default with the production failure threshold.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Add unit tests for QueryPlan and stored-procedure thin-client routing

Covers 5 new scenarios in ThinClientRoutingGateTests:
- ExecuteStoredProcedure on a StoredProcedure resource routes to thin client
- Non-execute StoredProcedure ops (Create) route to Gateway V1
- OperationType.QueryPlan routes to thin client
- QueryPlan returns false when probe is unhealthy
- ExecuteStoredProcedure returns false when probe is unhealthy

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Remove endpoint-probe content from QueryPlan PR branch

Reverts the merge of jeet1995/thin-client-probe-flow that was brought
into this PR earlier, while keeping the branch up-to-date with
upstream/main. Net diff vs upstream/main is now only the QueryPlan and
StoredProcedure Gateway V2 routing changes.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Propagate hybrid query diagnostics

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Advertise CountIf query feature

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address query plan review feedback

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Share thin client test property setup

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Clean up query range conversion state

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Thread DocumentCollection into readMany query-plan validation

readManyByPartitionKeys validates any caller-supplied custom query by
fetching a query plan and asserting it is single-partition and
non-hybrid. Until now that validation called fetchQueryPlanForValidation
with no DocumentCollection, so the plan request was pinned to Gateway V1
(useGatewayMode = (partitionKeyDefinition == null)).

Thread the container's DocumentCollection from
RxDocumentClientImpl.validateCustomQueryForReadManyByPartitionKeys ->
DocumentQueryExecutionContextFactory.fetchQueryPlanForValidation so
QueryPlanRetriever has the PartitionKeyDefinition needed to convert
PartitionKeyInternal-formatted queryRanges from the thin client (Gateway
V2) proxy into the EPK-hex Range<String> entries the query pipeline
consumes. With this wiring, the validation query plan goes to the thin
client when the client is configured for it, and remains on Gateway V1
otherwise. No behavior change on the non-thin-client path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Cover ReadManyByPartitionKey QueryPlan bifurcation with unit + E2E tests

- Add ReadManyByPartitionKeyQueryPlanRoutingTest unit tests that pin the useGatewayMode gate in QueryPlanRetriever: gateway mode when DocumentCollection is null and partitioned mode when a PartitionKeyDefinition is present.

- Add three readManyByPartitionKeys E2E tests to ThinClientQueryE2ETest that exercise the validation QueryPlan path through Direct TCP (baseline) and Gateway V2 (thin client), covering no-custom-query, projection+filter, and parameterized variants. Each thin-client diagnostics page is asserted to use the :10250 endpoint.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Skip RNTBD QueryPlan frames when capturing V2 feed requests

QueryPlan requests intentionally carry no RCS/CL headers (matches the V1
HTTP behavior). When the V2 thin-client routes the QueryPlan precursor
through the same :10250 endpoint as the data query, the spy must skip
the QueryPlan frame so the assertion checks the actual data-query frame.

This mirrors the IS_QUERY_PLAN_REQUEST filter on the V1 path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix thin-client QueryPlan error returning statusCode 0 instead of 400

The thin-client/Gateway-V2 proxy returns a non-2xx response with a raw,
non-JSON, NUL-padded error body for invalid-syntax queries. In
RxGatewayStoreModel.validateOrThrow, new CosmosError(body) attempted to
parse that body as JSON and threw IllegalArgumentException, which escaped
the method before the existing status-carrying throw could run. Upstream
then wrapped it as statusCode 0.

Wrap the CosmosError(body) construction in a narrow try/catch
(IllegalArgumentException) and fall back to the non-parsing
CosmosError(errorCode, message) constructor with a sanitized body. The
existing throw now fires with the correct status (400) and the proxy
error text is preserved.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Harden ThinClientQueryE2ETest assertions (F1-F5)

F1: strict 400 + unconditional thin-client endpoint check on invalid query, locking the statusCode-0->400 fix.
F2: ordered-vs-sorted-set document-ID comparison (ORDER BY sequence-compared, others set-compared).
F3: ID-set equality + no-duplicate check across drained continuation pages.
F4: numeric-tolerance (1e-6) comparison for scalar and GROUP BY aggregates to avoid float-formatting false mismatches.
F5: validated vector/full-text/hybrid queries match Direct vs thin-client end-to-end through the proxy.

Validated live via -Pthinclient: 84 tests, 0 failures.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Replace thin-client test matrix with reverse-engineered E2E test spec

Rewrite THINCLIENT_TEST_MATRIX.md as a reviewable test-design specification reverse-engineered from the committed ThinClientQueryE2ETest code (84 tests). Documents the differential-testing oracle (Direct :443 baseline vs thin client :10250 SUT), the data-model fixture, every assertion contract (endpoint provenance, ordered-vs-unordered ID equality, scalar/GROUP BY tolerance), the full 84-test matrix, the F1-F5 hardened special cases (continuation draining, invalid-query 400, vector/FTS/hybrid ranking, readMany validation path), advertised-feature coverage, and known gaps (CountIf/DCount/MultipleOrderBy) for reviewer sign-off.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Apply Aditya Sarpotdar review feedback to thin-client E2E tests

Executes actionable review feedback on ThinClientQueryE2ETest plus two
harness fixes surfaced during live validation against thin-client-multi-region-ci.

ThinClientQueryE2ETest:
- Strict ordering: ORDER BY results validated with isStrictlyOrdered across
  the board (stricter-by-default, per reviewer guidance) instead of set parity.
- Add testMultipleOrderBy() with a composite-index container to cover the
  MultipleOrderBy query feature.
- Add testDCount() using the canonical Cosmos DCount idiom
  (COUNT over a DISTINCT VALUE subquery); SQL-standard COUNT(DISTINCT ...) is
  not valid Cosmos SQL grammar.
- Reword multi-EPK-range comment/javadoc to reflect emulator/backend reality
  (multiple metadata ranges served by a single backend partition; SDK routing
  and query pipeline still exercised).

Harness fixes:
- TestSuiteBase: add wait-and-poll utility waitForCollectionToBeAvailableToRead
  (predicate on NotFound/substatus 1013) to deflake "Collection is not yet
  available for read" on freshly created containers; update call sites in
  OrderbyDocumentQueryTest, NonStreamingOrderByQueryVectorSearchTest,
  QueryValidationTests, ReadFeedCollectionsTest.
- SinglePartitionDocumentQueryTest: make the processMessage Mockito assertion
  mode-aware (thin-client routes query plan + query => times(2)).

Validated live against thin-client-multi-region-ci: 86/86 green.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address review feedback: opt-in thin-client QueryPlan routing + kill-switch, stop advertising CountIf, remove PR-reference comments

* Address review feedback: broaden gateway error fallback to ClassCastException, drop response ObjectNode mutation, fail-fast on empty queryRanges, align readMany routing test with opt-in

* Make thin-client QueryPlan no-EPK-headers contract explicit in wrapInHttpRequest

* Fix missing StandardCharsets import in RxGatewayStoreModelTest

* Test: add QueryOracle-derived LIKE/scalar-expression thin-client parity tests; drain only first page in OperationPolicies readAllItems to avoid full-container enumeration timeout

* Test: verify a cached proxy-generated query plan still executes correctly through the thin client

* Test: cross-partition tests use a dedicated multi-physical-partition fixture and assert >1 partition key range contacted; enforce ordered full-row comparison when id is not projected; clearer naming

* Add CHANGELOG entry for QueryPlan request routing to Gateway V2

* Call out Execute Stored Procedure support in CHANGELOG entry

* Fix flaky region/timing-sensitive fault-injection E2E tests

CustomerWorkflowPartitionLevelCircuitBreakerTest and CustomerWorkflowHighE2ETimeoutTest: bump the ThresholdBasedAvailabilityStrategy threshold from 100ms to 300ms so the cold primary-region connection can dispatch the request and record the injected fault before the availability-strategy hedge to a warm secondary region wins the race and cancels it (which previously yielded 0 fault hits).

PerPartitionCircuitBreakerE2ETests: before asserting short-circuit routing, wait until the first preferred region is actually marked Unavailable, removing the race between the failure-counter estimate and the asynchronous routing exclusion. Scoped to single-faulty-region configs (expectedRegionCountWithFailures == 1) and run once so fault-in-all-regions configs (which can never short-circuit) do not hit the method timeout.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Remove test AAD auth toggle; move thin-client test matrix to docs/ and align with ThinClientQueryE2ETest

* Add cross-partition aggregate / GROUP BY parity tests for thin-client query

Adds 7 thin-client E2E tests covering cross-partition aggregate merge
(COUNT/SUM/AVG/MIN/MAX) and GROUP BY, plus two helpers that assert
Direct vs thin-client value parity and verify the query genuinely
fanned out (distinctPartitionKeyRangesContacted > 1), exercising the
cross-partition MERGE path the single-partition aggregate tests miss.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix thin-client MULTI_HASH prefix EPK over-span; add QueryPlan parity tests

For a partial (prefix) hierarchical partition key, the thin-client / Gateway V2
path previously sent only StartEpkHash/EndEpkHash routing headers and no
doc-level EPK filter, so the proxy resolved the request to the owning physical
partition and returned every co-located document (e.g. 18 instead of 6).

ThinClientStoreModel.wrapInHttpRequest now sets READ_FEED_KEY_TYPE=
EffectivePartitionKeyRange, START_EPK and END_EPK on the request headers before
RntbdRequest.from() so RntbdRequestHeaders.addStartAndEndKeys serializes the
prefix EPK sub-range as the backend doc-level filter (hex string as UTF-8 bytes),
mirroring the Direct/RNTBD FeedRangeEpkImpl path. StartEpkHash/EndEpkHash routing
headers are retained.

Adds ThinClientQueryE2ETest coverage: QueryPlan parity across sources and a
hierarchical prefix half-open range test asserting thin-client matches Direct TCP.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Enforce QueryPlan thin-client routing invariant and harden container-readiness probe for thin-client-default-on

- Remove the QueryPlan exemption in assertThinClientEndpointUsed: when ThinClient is
  opted in with HTTP/2 enabled, QueryPlan calls MUST also route to Gateway V2 (:10250),
  matching production default DEFAULT_THINCLIENT_QUERY_PLAN_ENABLED=true.
- Restore public static visibility on ThinClientTestBase.assertThinClientEndpointUsed /
  assertGatewayEndpointUsed so cross-package static imports (CosmosMultiHashTest) compile.
- Add OWNER_RESOURCE_NOT_EXISTS (subStatus 1003) to the retryable NOTFOUND allowlist in
  TestSuiteBase.isRetryableCollectionReadinessFailure. With thin-client default-on the
  beforeSuite readability probe (a QueryPlan) routes through GW V2; a freshly-created
  container not yet propagated to a secondary region returns 404/1003, which must be
  treated as a transient readiness failure instead of cascade-skipping the suite.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* Add kill-switch flip validation for QueryPlan thin-client routing

Validate the COSMOS.THINCLIENT_QUERY_PLAN_ENABLED sysprop/env kill-switch
in both directions:
- Unit: new case in ReadManyByPartitionKeyQueryPlanRoutingTest asserts
  QueryPlan is forced to Gateway V1 (useGatewayMode=true) when the flag
  is disabled, even with thin-client eligible.
- E2E: make TestSuiteBase.assertThinClientEndpointUsed flag-aware so
  QueryPlan endpoint assertions flip based on the live-read flag while
  data operations continue to route to the thin-client endpoint.

Both scenarios verified green (unit 3/3; E2E 112/112, oracle:true for
flag ON and flag OFF).

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* Restore ambient COSMOS.THINCLIENT_ENABLED in ThinClientTestBase teardown

ThinClientTestBase teardown unconditionally cleared the
COSMOS.THINCLIENT_ENABLED system property. Since Configs.isThinClientEnabled()
reads the property lazily per client build, a ThinClientTestBase subclass
running before a property-dependent class inherited from main (e.g.
CosmosMultiHashTest, ThinClientQueryE2ETest) in the same JVM wiped the ambient
-DCOSMOS.THINCLIENT_ENABLED=true flag supplied by the thin-client CI lane.
Later clients then built with thin client disabled and routed to :443,
breaking the :10250 thin-client endpoint assertions.

Save the prior value on enable and restore it on teardown (clear only when it
was originally absent), mirroring CosmosConsistencyOverrideValidationTest's
restoreSystemProperty pattern. Helpers changed from static to instance so the
saved value is per-instance, avoiding races across TestNG factory instances.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* Remove redundant COSMOS.THINCLIENT_ENABLED mutation from ThinClientTestBase; revert workflow tests to main

ThinClientTestBase no longer sets/clears COSMOS.THINCLIENT_ENABLED per class. The 'thinclient' Maven lane already sets THINCLIENT_ENABLED=true and HTTP2_ENABLED=true JVM-wide via failsafe systemPropertyVariables, and all thin classes run only in that lane (groups={thinclient}). The prior per-class clearProperty in @afterclass wiped the ambient -D for later byte-identical-to-main classes in the same JVM, mis-routing requests to :443. Relying on the lane flag (matching ThinClientQueryE2ETest) removes that footgun.

Also revert CustomerWorkflowHighE2ETimeoutTest and CustomerWorkflowPartitionLevelCircuitBreakerTest to upstream/main (availability-strategy threshold stays 100ms), since this PR requires no changes there.

Verified locally (AZURE_TEST_MODE=LIVE, -Pthinclient,query,consistency-overrides): 234 passed, 0 failures/skips; QueryPlan routes to :10250.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Abhijeet Mohanty <copilot@noreply.local>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants